//
// Created by ECNU on 2024/2/28.
//

#ifndef NANO_FAISS_RESULTHANDLER_H
#define NANO_FAISS_RESULTHANDLER_H

#include <nanofaiss/utils/Heap.h>

namespace faiss {

template <class C>
struct BlockResultHandler {
    size_t nq; /// number of query
    size_t i0 = 0, i1 = 0; /// query range [i0, i1)

    explicit BlockResultHandler(size_t nq): nq(nq) {}

    /// start for search range between s and e
    virtual void begin_multiple(size_t s, size_t e) {
        i0 = s;
        i1 = e;
    }

    /// handle query [i0, i1) for base [j0, j1)
    virtual void add_results(size_t j0, size_t j1, const typename C::T*) {}

    /// end for the search
    virtual void end_multiple() {}


    virtual ~BlockResultHandler() {}

};


/// handle for a single query
template<class C>
struct ResultHandler {
    typename C::T threshold = 0; /// if not reach to the threshold, then not add to the result

    /// handle the single query, pure virtual function need to be implemented
    virtual bool add_result(typename C::T dis, typename C::TI idx) = 0;

    virtual ~ResultHandler() {}
};


/// Heap based result handler
template<class C>
struct HeapBlockResultHandler: BlockResultHandler<C> {
    using T = typename C::T;
    using TI = typename C::TI;
    using BlockResultHandler<C>::i0;
    using BlockResultHandler<C>::i1;

    T* heap_distances_tab;
    TI* heap_ids_tab;

    size_t k;

    HeapBlockResultHandler(
            size_t nq,
            T* heap_distances_tab,
            TI* heap_ids_tab,
            size_t k)
            : BlockResultHandler<C>(nq),
              heap_distances_tab(heap_distances_tab),
              heap_ids_tab(heap_ids_tab),
              k(k) {}

    /// API for 1 result at a time
    struct SingleResultHandler : ResultHandler<C> {
        HeapBlockResultHandler &hr; //! not need type name C, because inside the Class
        using ResultHandler<C>::threshold;
        size_t k;

        T* heap_distances;
        TI* heap_ids;

        explicit SingleResultHandler(HeapBlockResultHandler &hr)
                : hr(hr), k(hr.k) {}

        /// begin results for query #i
        void begin(size_t i) {
            heap_distances = hr.heap_distances_tab + i * k;
            heap_ids = hr.heap_ids_tab + i * k;
            heap_heapify<C>(k, heap_distances, heap_ids);
            threshold = heap_distances[0];
        }

        /// add the result to the heap and update the threshold
        bool add_result(T dis, TI idx) final {
            if (C::cmp(threshold, dis)) {
                heap_replace_top<C>(k, heap_distances, heap_ids, dis, idx);
                threshold = heap_distances[0];
                return true;
            }
            return false;
        }

        /// end the query and reorder the result
        void end() {
            heap_reorder<C>(k, heap_distances, heap_ids);
        }

    };

    /// API for multiple result

    /// begin
    void begin_multiple(size_t i0_, size_t i1_) final {
        i0 = i0_;
        i1 = i1_;
        for (size_t i = i0; i != i1; i++) {
            heap_heapify<C>(k, heap_distances_tab + i * k, heap_ids_tab + i * k);
        }
    }

    /// i0-i1 is the query, j0-j1 is the base id, so the dis_tab is (i0 - i1) * (j0 - j1)
    void add_results(size_t j0, size_t j1, const T* dis_tab) final {
#pragma omp parallel for
        for (size_t i = i0; i != i1; i++) {
            T* heap_dis = heap_distances_tab + i * k;
            TI* heap_ids = heap_ids_tab + i * k;
            const T *dis_tab_i = dis_tab + (j1 - j0) * (i - i0) - j0; // because the start of index is j0 in inner loop
            T thresh = heap_dis[0];
            for (size_t j = j0; j < j1; j++) {
                T dis = dis_tab_i[j];
                if (C::cmp(thresh, dis)) {
                    heap_replace_top<C>(k, heap_dis, heap_ids, dis, j);
                    thresh = heap_dis[0];
                }
            }
        }
    }


    void end_multiple() final {
#pragma omp for
        for (size_t i = i0; i < i1; i++)
            heap_reorder<C>(k, heap_distances_tab + i * k, heap_ids_tab + i * k);
    }
};
}

#endif // NANO_FAISS_RESULTHANDLER_H
